"""SQL Agent 节点模块 - 包含 SQL Agent 的所有节点功能。"""

from re import T
from langchain_core.messages import (
    AIMessage,
    HumanMessage,
    SystemMessage,
    trim_messages,
)
from langchain_core.runnables import RunnableConfig
from langchain_core.messages.tool import (
    InvalidToolCall,
    ToolCall,
    ToolCallChunk,
    ToolMessage,
    ToolMessageChunk,
    tool_call,
    tool_call_chunk,
)
from agents.states import SqlGraphAnnotation
from tools.database import list_tables_tool, get_schema_tool, run_query_tool
from tools.python_inter import python_inter, fig_inter

from utils.models import load_chat_model
from prompts.prompts import SQL_AGENT_SYSTEM, SQL_CHECKER_SYSTEM
from utils.logger import logger

import asyncio
from tenacity import retry, stop_after_attempt, wait_exponential


def list_tables(state: SqlGraphAnnotation, config: RunnableConfig = None) -> dict:
    """列出数据库中的所有表。"""
    # 使用标准的tool_call函数创建ToolCall对象
    tool_call_obj = tool_call(name="sql_db_list_tables", args={}, id="list_tables_001")

    # 创建包含tool_call的AIMessage
    tool_call_message = AIMessage(content="", tool_calls=[tool_call_obj])

    # 调用工具并获取ToolMessage响应
    tool_message = list_tables_tool.invoke(tool_call_obj)

    # logger.info(f"list_tables: {tool_message}")
    return {"messages": [tool_call_message, tool_message]}


def call_get_schema(state: SqlGraphAnnotation, config: RunnableConfig = None) -> dict:
    """调用获取数据库架构的工具。"""
    llm = load_chat_model()
    llm_with_tools = llm.bind_tools([get_schema_tool])
    # 修改：使用属性访问方式
    response = llm_with_tools.invoke(state.messages)
    # print(f"call_get_schema: {state.messages}")
    # logger.info(f"call_get_schema: {response}")
    return {"messages": [response]}


@retry(stop=stop_after_attempt(3), wait=wait_exponential(multiplier=1, min=4, max=10))
def generate_query(state: SqlGraphAnnotation, config: RunnableConfig = None) -> dict:
    """生成SQL查询。"""
    try:
        # 定义系统消息，专注于SQL查询规范
        system_message = SystemMessage(
            content="""你是一个SQL查询专家。在生成SQL查询时请遵循以下规范：
            1. 使用清晰的缩进和格式化，确保SQL查询可读性。
            2. 关键字使用大写（如SELECT, FROM, WHERE等）
            3. 表名和字段名使用小写
            4. 每个主要子句另起一行
            5. 使用有意义的别名
            6. 添加适当的注释说明查询逻辑
            7. 确保查询语法正确且高效
            8. 关于返回数据格式的问题，如果用户要求返回数据，请返回一个表格，并确保表格的格式正确，包含列名和数据
            请基于以上规范生成SQL查询。"""
        )

        llm = load_chat_model()
        llm_with_tools = llm.bind_tools([run_query_tool])
        # 修改：使用属性访问方式
        response = llm_with_tools.invoke([system_message] + state.messages)
        # print(f"generate_query: {state.messages}")
        # logger.info(f"generate_query: {response}")
        # logger.info(f"generate_query config: {config}")
        return {"messages": [response]}
    except asyncio.CancelledError:
        logger.warning("查询生成被取消")
        raise
    except Exception as e:
        logger.error(f"生成查询时出错: {str(e)}")
        raise


def check_query(state: SqlGraphAnnotation, config: RunnableConfig = None) -> dict:
    """检查并验证SQL查询。"""
    system_message = SystemMessage(content=SQL_CHECKER_SYSTEM)

    # 获取最后一个消息中的工具调用
    last_message = state.messages[-1]
    if hasattr(last_message, "tool_calls") and last_message.tool_calls:
        tool_call_obj = last_message.tool_calls[0]

        # 添加调试信息
        print(f"Tool call structure: {tool_call_obj}")
        # print(f"Tool call args: {tool_call_obj.get('args', {})}")
        # print(f"Available keys in args: {list(tool_call_obj.get('args', {}).keys())}")

        # 尝试不同的可能参数名
        args = tool_call_obj.get("args", {})
        query_text = None

        # 尝试常见的参数名
        for key in ["query", "sql", "sql_query", "input", "text"]:
            if key in args:
                query_text = args[key]
                print(f"Found query text with key '{key}': {query_text}")
                break

        if query_text is None:
            # 如果没有找到查询文本，使用第一个可用的值
            if args:
                query_text = list(args.values())[0]
                print(f"Using first available value: {query_text}")
            else:
                print("No args found in tool call")
                return {"messages": []}

        # 修复：创建ToolMessage时包含tool_call_id
        tool_message = ToolMessage(
            content=query_text, tool_call_id=tool_call_obj.get("id", "unknown_id")
        )

        llm = load_chat_model()
        llm_with_tools = llm.bind_tools([run_query_tool])
        response = llm_with_tools.invoke([system_message, tool_message])

        # 保持相同的ID以便正确更新
        if hasattr(last_message, "id"):
            response.id = last_message.id
        # logger.info(f"check_query: {response}")
        # logger.info(f"check_query config: {config}")

        return {"messages": [response]}

    return {"messages": []}


def trim_memory(state: SqlGraphAnnotation, config: RunnableConfig = None) -> dict:
    """截断短期记忆节点 - 在每次END之前调用，管理消息历史长度。

    该节点使用trim_messages函数来控制消息历史的长度，确保：
    1. 保持最近的对话历史
    2. 保留系统消息（如果存在）
    3. 确保对话历史以HumanMessage开始
    4. 控制token数量在合理范围内
    """
    try:
        # 如果消息列表为空或只有很少消息，直接返回
        if not state.messages or len(state.messages) <= 3:
            logger.info(f"消息数量较少({len(state.messages)})，跳过截断")
            return {"messages": state.messages}

        # 配置截断参数
        max_tokens = 4000  # 最大token数量，可根据需要调整

        # 使用trim_messages截断消息历史
        trimmed_messages = trim_messages(
            messages=state.messages,
            max_tokens=max_tokens,
            strategy="last",  # 保留最后的消息
            token_counter=load_chat_model(),  #
            start_on="human",  # 确保从HumanMessage开始
            end_on={"human", "tool"},  # 确保结束在HumanMessage或ToolMessage之间
            include_system=True,  # 确保最后一条消息被包含
        )

        # 记录截断信息
        original_count = len(state.messages)
        trimmed_count = len(trimmed_messages)

        if original_count != trimmed_count:
            logger.info(f"消息历史已截断: {original_count} -> {trimmed_count} 条消息")
        else:
            logger.info(f"消息历史无需截断，保持 {trimmed_count} 条消息")

        return {"messages": trimmed_messages}

    except Exception as e:
        logger.error(f"截断消息历史时出错: {str(e)}")
        # 如果截断失败，返回原始消息
        return {"messages": state.messages}
